import time
import re
from typing import List, Dict, Tuple, Generator
import numpy as np
import jieba
from collections import Counter
from dotenv import load_dotenv
import os

load_dotenv()

# OpenAI SDK用于DeepSeek API
from openai import OpenAI

# PDF处理和文本分割
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.schema import Document

# 重排序和交叉编码器
from sentence_transformers import CrossEncoder

# 关键词检索
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# FAISS优化
try:
    import faiss

    faiss_available = True
except ImportError:
    faiss_available = False
    print("FAISS不可用。请通过'pip install faiss-cpu'或'pip install faiss-gpu'安装")


class ContextEnhancer:
    """
    上下文增强器：用于增强多轮对话中的上下文相关性
    功能：
    1. 上下文相关性检测
    2. 历史对话重要性加权
    3. 指代消解
    4. 话题跟踪
    """

    def __init__(self, max_history_turns: int = 5):
        """
        初始化上下文增强器

        参数:
            max_history_turns: 历史对话保留的最大轮数
        """
        self.max_history_turns = max_history_turns
        self.topic_keywords = Counter()  # 用于跟踪对话主题的关键词计数器
        self.entities = {}  # 用于存储已识别的实体，用于指代消解
        self.llm_client = None  # 将在外部设置

    def set_llm_client(self, client):
        """设置LLM客户端，用于某些需要LLM辅助的处理"""
        self.llm_client = client

    def extract_keywords(self, text: str, top_n: int = 10) -> List[str]:
        """
        从文本中提取关键词

        参数:
            text: 输入文本
            top_n: 返回的关键词数量

        返回:
            关键词列表
        """
        # 使用jieba提取关键词
        words = jieba.analyse.extract_tags(text, topK=top_n)
        return words

    def update_topic_keywords(self, text: str, weight: float = 1.0):
        """
        更新话题关键词

        参数:
            text: 输入文本
            weight: 关键词权重
        """
        keywords = self.extract_keywords(text)
        for keyword in keywords:
            self.topic_keywords[keyword] += weight

    def get_topic_keywords(self, top_n: int = 10) -> List[str]:
        """
        获取主要话题关键词

        参数:
            top_n: 返回的关键词数量

        返回:
            关键词列表
        """
        return [kw for kw, _ in self.topic_keywords.most_common(top_n)]

    def extract_entities(self, text: str) -> List[str]:
        """
        从文本中提取实体

        参数:
            text: 输入文本

        返回:
            实体列表
        """
        # 简单实现：使用jieba分词，选择名词、人名等作为潜在实体
        # 在实际应用中，可以使用更复杂的命名实体识别方法
        words = jieba.posseg.cut(text)
        entities = []

        # 筛选可能的实体（名词、人名、地名等）
        for word, flag in words:
            if flag.startswith('n') and len(word) > 1:  # 名词
                entities.append(word)
            elif flag.startswith('nr'):  # 人名
                entities.append(word)
            elif flag.startswith('ns'):  # 地名
                entities.append(word)
            elif flag.startswith('nt'):  # 机构名
                entities.append(word)
            elif flag.startswith('nz'):  # 其他专名
                entities.append(word)

        return entities

    def update_entities(self, text: str, turn_index: int):
        """
        更新实体记录

        参数:
            text: 输入文本
            turn_index: 对话轮次索引
        """
        entities = self.extract_entities(text)
        for entity in entities:
            if entity not in self.entities:
                self.entities[entity] = []
            self.entities[entity].append(turn_index)

    def find_potential_references(self, text: str) -> List[Tuple[str, int]]:
        """
        查找可能的指代词

        参数:
            text: 输入文本

        返回:
            (指代词, 位置) 的列表
        """
        # 中文常见指代词
        reference_patterns = [
            r'[这那][个些]',  # 这个、那些等
            r'它[们的]?',  # 它、它们、它的
            r'他[们的]?',  # 他、他们、他的
            r'她[们的]?',  # 她、她们、她的
            r'该',  # 该
            r'此',  # 此
            r'上述',  # 上述
            r'前[述面]',  # 前述、前面
        ]

        references = []
        for pattern in reference_patterns:
            for match in re.finditer(pattern, text):
                references.append((match.group(), match.start()))

        return sorted(references, key=lambda x: x[1])

    def resolve_coreferences(self, text: str, history: List[Dict[str, str]]) -> str:
        """
        解析指代词，链接到之前提到的实体

        参数:
            text: 输入文本
            history: 对话历史

        返回:
            解析指代后的文本
        """
        if not history or not self.entities:
            return text

        # 查找可能的指代词
        references = self.find_potential_references(text)
        if not references:
            return text

        # 如果有指代词但没有实体记录，使用LLM进行指代消解
        if not self.entities and self.llm_client is not None:
            return self._resolve_with_llm(text, history)

        # 处理可能的实体替换
        # 这是一个简单实现，实际系统可能需要更复杂的逻辑
        result = text
        offset = 0  # 文本替换后的位置偏移

        for ref, pos in references:
            # 尝试找出最近提到的实体
            recent_entities = []
            for entity, turns in self.entities.items():
                if max(turns) < len(history):  # 确保实体出现在历史中
                    recent_entities.append((entity, max(turns)))

            if recent_entities:
                # 按最近提及排序
                recent_entities.sort(key=lambda x: x[1], reverse=True)
                replacement = recent_entities[0][0]  # 使用最近提到的实体

                # 替换文本中的指代词
                adjusted_pos = pos + offset
                result = result[:adjusted_pos] + replacement + result[adjusted_pos + len(ref):]
                offset += len(replacement) - len(ref)

        return result

    def _resolve_with_llm(self, text: str, history: List[Dict[str, str]]) -> str:
        """
        使用LLM进行指代消解

        参数:
            text: 输入文本
            history: 对话历史

        返回:
            解析指代后的文本
        """
        if not self.llm_client:
            return text

        # 构建提示
        history_text = "\n".join([
            f"{'用户' if msg['role'] == 'user' else '助手'}: {msg['content']}"
            for msg in history[-3:]  # 只取最近3轮对话
        ])

        prompt = f"""请帮我解析以下对话中的指代词（这个、那个、它、他们等），将这些指代词替换为它们所指的具体实体。
        不要改变原文的其他部分，只替换指代词。
        
        对话历史：
        {history_text}
        
        当前问题：
        {text}
        
        解析后的问题（只替换指代词）："""

        try:
            response = self.llm_client.chat.completions.create(
                model="deepseek-chat",  # 使用更轻量的模型
                messages=[{"role": "user", "content": prompt}],
                temperature=0,
                max_tokens=512
            )
            resolved_text = response.choices[0].message.content.strip()

            # 确保得到了合理的回复
            if len(resolved_text) > 0 and "解析后的问题" not in resolved_text:
                return resolved_text
        except Exception as e:
            print(f"LLM指代消解出错: {e}")

        return text

    def weight_history(self, history: List[Dict[str, str]]) -> List[Tuple[Dict[str, str], float]]:
        """
        对历史消息进行加权，越近的消息权重越高

        参数:
            history: 对话历史

        返回:
            带权重的对话历史
        """
        if not history:
            return []

        # 线性衰减权重
        weights = []
        for i in range(len(history)):
            # 用户消息权重略高于助手回复
            is_user = history[i]['role'] == 'user'
            position_weight = (i + 1) / len(history)  # 位置权重：越新的消息权重越高
            role_weight = 1.2 if is_user else 1.0  # 角色权重：用户消息权重稍高
            weights.append(position_weight * role_weight)

        # 归一化权重
        total_weight = sum(weights)
        weights = [w / total_weight for w in weights]

        return list(zip(history, weights))

    def enhance_query(self, query: str, history: List[Dict[str, str]]) -> str:
        """
        增强查询，结合对话历史和话题

        参数:
            query: 原始查询
            history: 对话历史

        返回:
            增强后的查询
        """
        if not history:
            # 首次查询，更新话题关键词
            self.update_topic_keywords(query, weight=2.0)  # 首次查询权重更高
            return query

        # 解析指代词
        resolved_query = self.resolve_coreferences(query, history)

        # 提取当前查询的关键词
        query_keywords = self.extract_keywords(resolved_query)

        # 更新话题关键词
        self.update_topic_keywords(resolved_query)

        # 加权历史消息
        weighted_history = self.weight_history(history)

        # 从历史中提取相关上下文
        context_parts = []
        for msg, weight in weighted_history:
            if msg['role'] == 'user':
                # 提取关键词
                msg_keywords = self.extract_keywords(msg['content'])

                # 计算与当前查询的关键词重叠
                overlap = set(query_keywords) & set(msg_keywords)
                if overlap:
                    # 有重叠的关键词，添加到上下文
                    context_parts.append(msg['content'])

        # 如果有相关历史，构建增强查询
        if context_parts:
            # 选择最近的2条相关历史
            recent_contexts = context_parts[-2:] if len(context_parts) > 2 else context_parts

            # 构建增强查询（不直接合并，而是作为参考）
            enhanced_query = f"{resolved_query} (参考上下文: {' '.join(recent_contexts)})"
            return enhanced_query

        return resolved_query


class OptimizedRAGSystem:
    """
    优化的RAG系统，集成多种技术以提高大型文档处理性能：
    1. 领域适应的嵌入模型
    2. 混合检索策略（TF-IDF + 向量检索）
    3. 重排序机制
    4. 向量量化技术
    5. 多轮对话支持
    6. 上下文相关性增强
    """

    def __init__(
            self,
            embedding_model_name: str = "shibing624/text2vec-base-chinese",
            reranker_model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v1",
            vector_weight: float = 0.7,
            use_quantization: bool = True,
            recall_size: int = 100,
            chunk_size: int = 500,
            chunk_overlap: int = 50,
            deepseek_api_key: str = None,
            deepseek_api_base: str = "https://api.deepseek.com/v1",
            max_history_turns: int = 5
    ):
        """
        初始化优化的RAG系统

        参数:
            embedding_model_name: 嵌入模型名称
            reranker_model_name: 重排序模型名称
            vector_weight: 向量检索权重
            use_quantization: 是否使用向量量化
            recall_size: 初始召回文档数量
            chunk_size: 文档块大小
            chunk_overlap: 块重叠大小
            deepseek_api_key: DeepSeek API密钥
            deepseek_api_base: DeepSeek API基础URL
            max_history_turns: 历史对话保留的最大轮数
        """
        self.vector_weight = vector_weight
        self.recall_size = recall_size
        self.use_quantization = use_quantization
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.max_history_turns = max_history_turns

        # 初始化嵌入模型
        print(f"加载嵌入模型: {embedding_model_name}")
        self.embeddings = HuggingFaceEmbeddings(
            model_name=embedding_model_name,
            model_kwargs={'device': 'cpu'},
            encode_kwargs={'normalize_embeddings': True}
        )

        # 初始化重排序模型
        print(f"加载重排序模型: {reranker_model_name}")
        self.reranker = CrossEncoder(reranker_model_name)

        # 使用OpenAI SDK初始化DeepSeek LLM
        self.client = OpenAI(
            api_key=deepseek_api_key,
            base_url=deepseek_api_base
        )

        # 初始化上下文增强器
        self.context_enhancer = ContextEnhancer(max_history_turns=max_history_turns)
        self.context_enhancer.set_llm_client(self.client)

        # 存储变量
        self.documents = []
        self.chunks = []
        self.vector_store = None
        self.keyword_retriever = None

        # 对话历史存储
        self.conversation_histories = {}  # 用于存储不同会话的对话历史的字典

    def get_or_create_conversation_history(self, session_id: str = "default") -> List[Dict[str, str]]:
        """
        获取或创建会话的对话历史

        参数:
            session_id: 会话标识符

        返回:
            对话历史消息列表
        """
        if session_id not in self.conversation_histories:
            self.conversation_histories[session_id] = []

        return self.conversation_histories[session_id]

    def add_to_conversation_history(self, session_id: str, role: str, content: str):
        """
        向对话历史添加消息

        参数:
            session_id: 会话标识符
            role: 消息角色（用户/助手）
            content: 消息内容
        """
        history = self.get_or_create_conversation_history(session_id)

        # 添加到历史记录
        history.append({"role": role, "content": content})

        # 如果是用户消息，更新上下文增强器的实体和话题
        if role == "user":
            turn_index = len(history) - 1
            self.context_enhancer.update_entities(content, turn_index)
            self.context_enhancer.update_topic_keywords(content)

        # 只保留最后max_history_turns轮用户-助手交流
        if len(history) > self.max_history_turns * 2:
            # 只保留最后max_history_turns轮交流
            history = history[-(self.max_history_turns * 2):]
            self.conversation_histories[session_id] = history

    def clear_conversation_history(self, session_id: str = "default"):
        """
        清除会话的对话历史

        参数:
            session_id: 会话标识符
        """
        if session_id in self.conversation_histories:
            self.conversation_histories[session_id] = []
            # 重置上下文增强器的状态
            self.context_enhancer = ContextEnhancer(max_history_turns=self.max_history_turns)
            self.context_enhancer.set_llm_client(self.client)

    def load_pdf(self, pdf_path: str) -> List[Document]:
        """
        加载PDF文档并分割成块

        参数:
            pdf_path: PDF文件路径

        返回:
            文档块列表
        """
        print(f"加载PDF文档: {pdf_path}")
        start_time = time.time()

        # 使用PyPDFLoader加载PDF
        loader = PyPDFLoader(pdf_path)
        documents = loader.load()

        # 文档分块
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap,
            separators=["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""]
        )
        chunks = text_splitter.split_documents(documents)

        print(f"PDF处理完成，{len(documents)}页，分割成{len(chunks)}个块")
        print(f"处理时间: {time.time() - start_time:.2f}秒")

        self.documents = documents
        self.chunks = chunks
        return chunks

    def build_retrieval_system(self):
        """构建检索系统，包括向量存储和关键词检索器"""
        if not self.chunks:
            raise ValueError("请先加载文档")

        start_time = time.time()

        # 构建向量存储
        print("构建向量存储...")
        texts = [chunk.page_content for chunk in self.chunks]
        metadatas = [chunk.metadata for chunk in self.chunks]

        if self.use_quantization and faiss_available:
            # 使用FAISS构建带有IVF-PQ量化的向量存储
            # 1. 首先获取向量嵌入
            embeddings_list = self.embeddings.embed_documents(texts)

            # 2. 创建FAISS索引（带量化）
            dimension = len(embeddings_list[0])
            self.vector_store = FAISS.from_embeddings(
                text_embeddings=list(zip(texts, embeddings_list)),
                embedding=self.embeddings,
                metadatas=metadatas
            )

            # 3. 应用IVF-PQ量化技术
            index = self.vector_store.index
            if len(texts) > 10000:  # 对大数据使用量化
                n_list = 1000  # 聚类中心数量
                m = 8  # 子量化器数量

                # 创建原始索引的副本
                orig_index = faiss.extract_index_ivf(index)

                # 构建新的IVF-PQ索引
                quantizer = faiss.IndexFlatIP(dimension)
                ivf_pq_index = faiss.IndexIVFPQ(quantizer, dimension, n_list, m, 8)
                ivf_pq_index.train(np.array(embeddings_list))

                # 添加向量
                for i, embedding in enumerate(embeddings_list):
                    ivf_pq_index.add_with_ids(np.array([embedding]), np.array([i]))

                # 替换原始索引
                self.vector_store.index = ivf_pq_index
                self.vector_store.index.nprobe = 10
                print(f"应用了IVF-PQ量化，内存使用减少约75%")

        else:
            # 常规FAISS索引
            self.vector_store = FAISS.from_texts(texts, self.embeddings, metadatas=metadatas)

        # 构建关键词检索器
        print("构建关键词检索器...")
        self.keyword_retriever = KeywordRetriever(texts, metadatas)

        print(f"检索系统构建完成，时间: {time.time() - start_time:.2f}秒")

    def retrieve(self, query: str, top_k: int = 5, history: List[Dict[str, str]] = None) -> List[Document]:
        """
        混合检索 + 重排序策略，考虑历史对话和关键词

        参数:
            query: 查询文本
            top_k: 返回文档数量
            history: 对话历史

        返回:
            检索到的文档列表
        """
        if not self.vector_store or not self.keyword_retriever:
            raise ValueError("请先构建检索系统")

        start_time = time.time()

        # 1. 使用上下文增强器增强查询
        if history:
            enhanced_query = self.context_enhancer.enhance_query(query, history)
            print(f"增强后的查询: {enhanced_query}")
        else:
            enhanced_query = query

        # 2. 向量检索
        vector_results = self.vector_store.similarity_search_with_score(
            enhanced_query, k=self.recall_size
        )

        # 3. 关键词检索
        keyword_results = self.keyword_retriever.retrieve(enhanced_query, top_k=self.recall_size)

        # 4. 混合结果
        id_to_score = {}
        id_to_doc = {}

        # 处理向量检索结果
        for doc, score in vector_results:
            doc_id = str(hash(doc.page_content))
            normalized_score = 1.0 - (score / 2.0)  # 将距离转换为相似度分数
            id_to_score[doc_id] = self.vector_weight * normalized_score
            id_to_doc[doc_id] = doc

        # 处理关键词检索结果
        for doc, score in keyword_results:
            doc_id = str(hash(doc["text"]))
            id_to_score[doc_id] = id_to_score.get(doc_id, 0) + (1 - self.vector_weight) * score
            if doc_id not in id_to_doc:
                langchain_doc = Document(
                    page_content=doc["text"],
                    metadata=doc["metadata"]
                )
                id_to_doc[doc_id] = langchain_doc

        # 5. 话题相关性提升
        if history and self.context_enhancer.topic_keywords:
            topic_keywords = self.context_enhancer.get_topic_keywords()

            # 为包含主题关键词的文档增加分数
            for doc_id, doc in id_to_doc.items():
                for keyword in topic_keywords:
                    if keyword in doc.page_content:
                        # 为包含主题关键词的文档增加权重
                        id_to_score[doc_id] = id_to_score.get(doc_id, 0) + 0.1

        # 排序并选择前recall_size个结果
        sorted_ids = sorted(id_to_score.keys(), key=lambda k: id_to_score[k], reverse=True)[:self.recall_size]
        hybrid_results = [(id_to_doc[doc_id], id_to_score[doc_id]) for doc_id in sorted_ids]

        print(f"混合检索完成，时间: {time.time() - start_time:.2f}秒")

        # 6. 重排序，加入对话历史因素
        if len(hybrid_results) > top_k:
            rerank_start = time.time()
            print(f"开始重排序，对{len(hybrid_results)}个候选项进行排序...")

            # 准备重排序输入
            if history:
                # 将历史信息添加到查询中，以考虑历史对话
                recent_history = history[-4:]  # 只使用最近的几轮对话
                history_text = " ".join([msg["content"] for msg in recent_history])
                context_query = f"{enhanced_query} {history_text}"
            else:
                context_query = enhanced_query

            pairs = [(context_query, doc.page_content) for doc, _ in hybrid_results]
            rerank_scores = self.reranker.predict(pairs)

            # 重新排序
            reranked_results = sorted(
                [(doc, score) for (doc, _), score in zip(hybrid_results, rerank_scores)],
                key=lambda x: x[1],
                reverse=True
            )

            # 取前top_k个结果
            top_results = [doc for doc, _ in reranked_results[:top_k]]
            print(f"重排序完成，时间: {time.time() - rerank_start:.2f}秒")
        else:
            top_results = [doc for doc, _ in hybrid_results[:top_k]]

        print(f"总检索时间: {time.time() - start_time:.2f}秒，返回{len(top_results)}个文档")
        return top_results

    def answer_question(self, query: str, top_k: int = 5, session_id: str = "default") -> str:
        """
        使用检索到的文档和对话历史回答问题

        参数:
            query: 查询文本
            top_k: 检索文档数量
            session_id: 会话标识符

        返回:
            回答文本
        """
        # 获取对话历史
        history = self.get_or_create_conversation_history(session_id)

        # 检索相关文档，考虑历史对话
        relevant_docs = self.retrieve(query, top_k=top_k, history=history)

        # 从检索到的文档创建上下文
        context = "\n\n".join([doc.page_content for doc in relevant_docs])

        # 创建带有文档上下文的系统消息
        system_message = {"role": "system", "content": f"""你是一个专业的助手，擅长根据提供的文档内容回答问题。请始终用中文回答问题。
        
        文档内容:
        {context}
        
        用户当前的问题是: {query}
        
        使用此文档内容回答用户的问题。如果文档内容不足以回答问题，可以使用对话历史或表明你无法回答。
        在回答时，注意保持与对话历史的一致性，避免自相矛盾。
        """}

        # 将用户查询添加到历史记录
        self.add_to_conversation_history(session_id, "user", query)

        # 创建API调用的消息
        api_messages = [system_message] + history

        # 使用DeepSeek生成回答
        response = self.client.chat.completions.create(
            model="deepseek-reasoner",
            messages=api_messages,
            temperature=0,
            max_tokens=1024
        )

        answer = response.choices[0].message.content

        # 将助手响应添加到历史记录
        self.add_to_conversation_history(session_id, "assistant", answer)

        return answer

    def answer_question_stream(self, query: str, top_k: int = 5, session_id: str = "default") -> Generator[
        str, None, None]:
        """
        使用对话历史流式生成回答

        参数:
            query: 查询文本
            top_k: 检索文档数量
            session_id: 会话标识符

        返回:
            生成回答文本块的生成器
        """
        # 获取对话历史
        history = self.get_or_create_conversation_history(session_id)

        # 检索相关文档，考虑历史对话
        relevant_docs = self.retrieve(query, top_k=top_k, history=history)

        # 从检索到的文档创建上下文
        context = "\n\n".join([doc.page_content for doc in relevant_docs])

        # 创建带有文档上下文的系统消息
        system_message = {"role": "system", "content": f"""你是一个专业的助手，擅长根据提供的文档内容回答问题。请始终用中文回答问题。
        
        文档内容:
        {context}
        
        用户当前的问题是: {query}
        
        使用此文档内容回答用户的问题。如果文档内容不足以回答问题，可以使用对话历史或表明你无法回答。
        在回答时，注意保持与对话历史的一致性，避免自相矛盾。
        """}

        # 将用户查询添加到历史记录
        self.add_to_conversation_history(session_id, "user", query)

        # 创建API调用的消息
        api_messages = [system_message] + history

        # 使用DeepSeek流式生成回答
        stream = self.client.chat.completions.create(
            model="deepseek-reasoner",
            messages=api_messages,
            temperature=0,
            max_tokens=1024,
            stream=True
        )

        # 收集完整响应以添加到历史记录
        full_response = ""

        for chunk in stream:
            if chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                full_response += content
                yield content

        # 将助手响应添加到历史记录
        self.add_to_conversation_history(session_id, "assistant", full_response)

    def get_conversation_topics(self, session_id: str = "default") -> List[str]:
        """
        获取当前对话的主要话题

        参数:
            session_id: 会话标识符

        返回:
            话题关键词列表
        """
        return self.context_enhancer.get_topic_keywords()


class KeywordRetriever:
    """基于TF-IDF的关键词检索器"""

    def __init__(self, texts: List[str], metadatas: List[Dict[str, any]]):
        """
        初始化关键词检索器

        参数:
            texts: 文本列表
            metadatas: 元数据列表
        """
        self.texts = texts
        self.metadatas = metadatas

        # 使用jieba分词器
        self.vectorizer = TfidfVectorizer(
            tokenizer=lambda x: list(jieba.cut(x)),
            lowercase=True
        )

        self.tfidf_matrix = self.vectorizer.fit_transform(texts)

    def retrieve(self, query: str, top_k: int = 5) -> List[Tuple[Dict, float]]:
        """
        检索最相关的文档

        参数:
            query: 查询文本
            top_k: 返回文档数量

        返回:
            检索到的文档和分数列表
        """
        start_time = time.time()

        # 向量化查询
        query_vector = self.vectorizer.transform([query])

        # 计算相似度
        similarities = cosine_similarity(query_vector, self.tfidf_matrix)[0]

        # 选择前top_k个结果
        top_indices = similarities.argsort()[::-1][:top_k]
        results = [
            ({"text": self.texts[idx], "metadata": self.metadatas[idx]}, similarities[idx])
            for idx in top_indices
        ]

        print(f"关键词检索时间: {time.time() - start_time:.2f}秒")
        return results


def interactive_chat():
    """交互式命令行聊天函数"""

    # 获取API密钥（可以通过环境变量或直接输入）
    api_key = os.environ.get("DEEPSEEK_API_KEY", None)
    if not api_key:
        api_key = input("请输入您的DeepSeek API密钥: ")

    # 获取PDF路径
    pdf_path = input("请输入PDF文件路径: ")

    # 初始化RAG系统
    print("\n初始化RAG系统...")
    rag_system = OptimizedRAGSystem(
        embedding_model_name="shibing624/text2vec-base-chinese",
        reranker_model_name="cross-encoder/ms-marco-MiniLM-L-6-v1",
        vector_weight=0.7,
        use_quantization=True,
        recall_size=100,
        chunk_size=500,
        chunk_overlap=50,
        deepseek_api_key=api_key,
        max_history_turns=5
    )

    # 加载PDF文档
    print("\n加载PDF文档...")
    rag_system.load_pdf(pdf_path)

    # 构建检索系统
    print("\n构建检索系统...")
    rag_system.build_retrieval_system()

    # 创建会话ID
    session_id = f"session_{int(time.time())}"

    print("\n=== 开始交互式对话 ===")
    print("输入 'exit' 或 'quit' 结束对话")
    print("输入 'clear' 清除对话历史")
    print("输入 'stream' 切换到流式输出模式")
    print("输入 'topics' 显示当前对话主题")

    # 默认不使用流式输出
    use_stream = False

    # 交互式对话循环
    while True:
        # 获取用户输入
        query = input("\n用户: ")

        # 检查特殊命令
        if query.lower() in ['exit', 'quit']:
            print("再见！")
            break
        elif query.lower() == 'clear':
            rag_system.clear_conversation_history(session_id)
            print("已清除对话历史")
            continue
        elif query.lower() == 'stream':
            use_stream = not use_stream
            print(f"流式输出模式: {'开启' if use_stream else '关闭'}")
            continue
        elif query.lower() == 'topics':
            topics = rag_system.get_conversation_topics(session_id)
            print(f"当前对话主题关键词: {', '.join(topics)}")
            continue

        # 处理正常问题
        if use_stream:
            print("助手: ", end="", flush=True)
            for chunk in rag_system.answer_question_stream(query, top_k=5, session_id=session_id):
                print(chunk, end="", flush=True)
            print()  # 换行
        else:
            answer = rag_system.answer_question(query, top_k=5, session_id=session_id)
            print(f"助手: {answer}")


if __name__ == "__main__":
    # 导入jieba分析模块，用于关键词提取
    import jieba.analyse

    interactive_chat()
